from ultralytics import YOLO
import cv2
import argparse

# 创建 ArgumentParser 对象
parser = argparse.ArgumentParser(description='输入输出文件')

# 添加参数
parser.add_argument('-f', help='输入文件')
parser.add_argument('-d', help='输出文件', default='default value')


# Load a model
#model = YOLO("yolo11n.pt")

def model_train():
    # Train the model
    print("start to train")
    train_results = model.train(
        data="webfuture.yaml",  # path to dataset YAML
        epochs=50,  # number of training epochs
        imgsz=640,  # training image size
        device="cpu",  # device to run on, i.e. device=0 or device=0,1,2,3 or device=cpu
    )
    print(train_results)

def model_val():
    # Evaluate model performance on the validation set
    metrics = model.val()
    # print(metrics)

def model_export():
    # Export the model to ONNX format
    return  model.export(format="onnx")  # return path to exported model


if __name__ == '__main__':
    args = parser.parse_args()
    
    model = YOLO("yolo11m.pt")

    model_train()
    model_val()
    model_export()

# python yolov11-detect.py -f /Users/v/Documents/proj/testdata/origin_imgs/images_labelimg_yolo/IMG_20181201_181804.jpg -d /Users/v/Documents/proj/testdata/origin_imgs/r1-out.jpg

